import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional
import math
from dataclasses import dataclass
from collections import deque
import time
from numba import jit


@dataclass
class PlatParaInfo:
    """我方平台参数信息"""
    id: int
    longitude: float  # 经度
    latitude: float   # 纬度
    altitude: float   # 高度
    speed: float      # 速度 fSpeed
    roll: float       # 翻滚角
    azimuth: float    # 水平方位角 fDir
    elevation: float  # 仰角
    timestamp: float  # 时间戳

@dataclass
class RadarInfo:
    """雷达探测信息"""
    enemy_id: int     # 敌机ID
    longitude: float  # 经度
    latitude: float   # 纬度
    altitude: float   # 高度
    timestamp: float  # 时间戳

class CoordinateTransformer:
    """坐标转换类"""
    
    # WGS84椭球参数
    a = 6378137.0  # 长半轴
    f = 1/298.257223563  # 扁率
    b = a * (1 - f)  # 短半轴
    e2 = 1 - (b/a)**2  # 第一偏心率平方
    
    @staticmethod
   
    def llh_to_ecef(lat, lon, alt):
        """将经纬高转换为ECEF坐标"""
        lat_rad = np.radians(lat)
        lon_rad = np.radians(lon)
        
        N = CoordinateTransformer.a / np.sqrt(1 - CoordinateTransformer.e2 * np.sin(lat_rad)**2)
        
        x = (N + alt) * np.cos(lat_rad) * np.cos(lon_rad)
        y = (N + alt) * np.cos(lat_rad) * np.sin(lon_rad)
        z = (N * (1 - CoordinateTransformer.e2) + alt) * np.sin(lat_rad)
        
        return np.array([x, y, z])
    
    @staticmethod
    def ecef_to_enu(ecef, ref_lat, ref_lon, ref_alt):
        """将ECEF坐标转换为ENU坐标（以参考点为原点）"""
        ref_lat_rad = np.radians(ref_lat)
        ref_lon_rad = np.radians(ref_lon)
        
        # 计算旋转矩阵
        R = np.array([
            [-np.sin(ref_lon_rad), np.cos(ref_lon_rad), 0],
            [-np.sin(ref_lat_rad)*np.cos(ref_lon_rad), -np.sin(ref_lat_rad)*np.sin(ref_lon_rad), np.cos(ref_lat_rad)],
            [np.cos(ref_lat_rad)*np.cos(ref_lon_rad), np.cos(ref_lat_rad)*np.sin(ref_lon_rad), np.sin(ref_lat_rad)]
        ])
        
        # 参考点ECEF坐标
        ref_ecef = CoordinateTransformer.llh_to_ecef(ref_lat, ref_lon, ref_alt)
        
        # 相对坐标
        delta = ecef - ref_ecef
        
        # 转换为ENU
        enu = R @ delta
        return enu
    
    @staticmethod
    def llh_to_enu(lat, lon, alt, ref_lat, ref_lon, ref_alt):
        """直接将经纬高转换为ENU坐标"""
        ecef = CoordinateTransformer.llh_to_ecef(lat, lon, alt)
        return CoordinateTransformer.ecef_to_enu(ecef, ref_lat, ref_lon, ref_alt)


class HighPrecisionStateEstimator:
    """高精度状态估计器（基于卡尔曼滤波）"""
    
    def __init__(self, state_dim=9, obs_dim=3):
        self.state_dim = state_dim  # [x, y, z, vx, vy, vz, ax, ay, az]
        self.obs_dim = obs_dim
        
        # 初始化状态向量和协方差矩阵
        self.x = np.zeros(state_dim)
        self.P = np.eye(state_dim) * 1000  # 初始不确定性较大
        
        # 过程噪声协方差
        self.Q = np.eye(state_dim) * 0.1
        
        # 观测矩阵
        self.H = np.zeros((obs_dim, state_dim))
        self.H[:3, :3] = np.eye(3)
        
        # 观测噪声协方差
        self.R = np.eye(obs_dim) * 10
        
        self.last_timestamp = None
        
    def predict(self, current_time):
        """预测步骤"""
        if self.last_timestamp is None:
            self.last_timestamp = current_time
            return
            
        dt = current_time - self.last_timestamp
        self.last_timestamp = current_time
        
        if dt <= 0:
            return
            
        # 状态转移矩阵（三阶运动模型）
        F = self._compute_state_transition_matrix(dt)
        
        # 预测状态
        self.x = F @ self.x
        
        # 预测协方差
        self.P = F @ self.P @ F.T + self.Q
        
    def update(self, measurement, current_time):
        """更新步骤"""
        if self.last_timestamp is None:
            self.last_timestamp = current_time
            
        dt = current_time - self.last_timestamp
        if dt > 0:
            self.predict(current_time)
        
        # 计算卡尔曼增益
        S = self.H @ self.P @ self.H.T + self.R
        K = self.P @ self.H.T @ np.linalg.inv(S)
        
        # 更新状态
        y = measurement - self.H @ self.x  # 新息
        self.x = self.x + K @ y
        
        # 更新协方差
        I = np.eye(self.state_dim)
        self.P = (I - K @ self.H) @ self.P
        
    def _compute_state_transition_matrix(self, dt):
        """计算状态转移矩阵"""
        F = np.eye(self.state_dim)
    
        # 位置更新
        F[0, 3] = dt  # x += vx * dt
        F[1, 4] = dt  # y += vy * dt
        F[2, 5] = dt  # z += vz * dt
        
        F[0, 6] = 0.5 * dt**2  # x += 0.5 * ax * dt^2
        F[1, 7] = 0.5 * dt**2  # y += 0.5 * ay * dt^2
        F[2, 8] = 0.5 * dt**2  # z += 0.5 * az * dt^2
        
        # 速度更新
        F[3, 6] = dt  # vx += ax * dt
        F[4, 7] = dt  # vy += ay * dt
        F[5, 8] = dt  # vz += az * dt
        
        return F
    
    def get_state(self):
        """获取当前状态估计"""
        return {
            'position': self.x[:3],
            'velocity': self.x[3:6],
            'acceleration': self.x[6:9],
            'covariance': self.P
        }
    

class TimeAwareThreatEvaluator:
    """时间感知的高精度威胁评估器"""
    
    def __init__(self):
        self.enemy_tracks = {}  # 敌机轨迹记录
        self.friendly_platforms = {}  # 我方平台信息
        self.threat_history = {}  # 威胁历史记录
        
        # 威胁参数配置
        self.threat_params = {
            'distance': {'k_d': 0.001, 'd_max': 50000},
            'angle': {'max_angle': np.pi},
            'speed': {'v_max': 300, 'v_closure_max': 100},
            'behavior': {'maneuver_threshold': 2.0, 'curvature_threshold': 0.01}
        }
        
        # 动态权重参数
        self.weights = {
            'base': {'distance': 0.35, 'angle': 0.25, 'speed': 0.20, 'behavior': 0.20},
            'adaptive_gain': 0.1  # 自适应调整增益
        }
        
    def update_friendly_info(self, plat_info: PlatParaInfo):
        """更新我方平台信息"""
        self.friendly_platforms[plat_info.id] = plat_info
        
    def update_radar_info(self, radar_info: List[RadarInfo]):
        """更新雷达信息并处理敌机轨迹"""
        current_time = time.time()
        
        for radar_point in radar_info:
            enemy_id = radar_point.enemy_id
            
            if enemy_id not in self.enemy_tracks:
                # 初始化敌机跟踪
                self.enemy_tracks[enemy_id] = {
                    'estimator': HighPrecisionStateEstimator(),
                    'position_history': deque(maxlen=50),
                    'time_history': deque(maxlen=50),
                    'threat_history': deque(maxlen=20)
                }
            
            track = self.enemy_tracks[enemy_id]
            
            # 将雷达点转换为ENU坐标（以第一个我方平台为参考）
            if self.friendly_platforms:
                ref_platform = next(iter(self.friendly_platforms.values()))
                enu_pos = CoordinateTransformer.llh_to_enu(
                    radar_point.latitude, radar_point.longitude, radar_point.altitude,
                    ref_platform.latitude, ref_platform.longitude, ref_platform.altitude
                )
                
                # 更新状态估计器
                track['estimator'].update(enu_pos, radar_point.timestamp)
                track['position_history'].append(enu_pos)
                track['time_history'].append(radar_point.timestamp)
    
    def calculate_high_precision_closure_speed(self, friendly_pos, friendly_vel, enemy_pos, enemy_vel, dt):
        """计算高精度接近速度"""
        # 相对位置和速度
        relative_pos = enemy_pos - friendly_pos
        relative_vel = enemy_vel - friendly_vel
        
        distance = np.linalg.norm(relative_pos)
        if distance < 1e-6:
            return 0.0
        
        # 基础接近速度
        v_closure_basic = -np.dot(relative_vel, relative_pos) / distance
        
        # 考虑加速度的高阶修正
        if len(self.enemy_tracks) > 1:
            # 计算相对加速度（需要历史数据）
            acc_correction = self._calculate_acceleration_correction(friendly_pos, enemy_pos, dt)
            v_closure_enhanced = v_closure_basic + 0.5 * acc_correction * dt
        else:
            v_closure_enhanced = v_closure_basic
            
        return max(v_closure_enhanced, 0)  # 确保非负
    
    def _calculate_acceleration_correction(self, friendly_pos, enemy_pos, dt):
        """计算加速度修正项"""
        # 简化的加速度估计（实际应用中可以使用更复杂的方法）
        if dt > 0:
            # 基于位置变化的加速度估计
            return 0.0  # 简化实现
        return 0.0
    
    def calculate_dynamic_weights(self, friendly_id, enemy_id, current_time):
        """计算动态权重"""
        base_weights = self.weights['base'].copy()
        adaptive_gain = self.weights['adaptive_gain']
        
        if enemy_id not in self.enemy_tracks:
            return base_weights
        
        track = self.enemy_tracks[enemy_id]
        
        if len(track['threat_history']) < 2:
            return base_weights
        
        # 计算各威胁分量的时间变化率
        threat_variations = self._calculate_threat_variations(track['threat_history'])
        
        # 根据变化率调整权重（变化快的分量权重增加）
        for component in base_weights.keys():
            if component in threat_variations:
                variation = abs(threat_variations[component])
                base_weights[component] *= (1 + adaptive_gain * variation)
        
        # 归一化权重
        total_weight = sum(base_weights.values())
        normalized_weights = {k: v/total_weight for k, v in base_weights.items()}
        
        return normalized_weights
    
    def _calculate_threat_variations(self, threat_history):
        """计算威胁分量变化率"""
        if len(threat_history) < 2:
            return {}
        
        # 提取最近的两个威胁评估
        recent_threats = list(threat_history)[-2:]
        variations = {}
        
        for component in recent_threats[0]['components'].keys():
            if component in recent_threats[1]['components']:
                delta = recent_threats[1]['components'][component] - recent_threats[0]['components'][component]
                # 假设时间间隔为1秒（简化）
                variations[component] = delta
        
        return variations
    
    def evaluate_threat(self, friendly_id, enemy_id):
        """评估单个敌机对单个我机的威胁"""
        if friendly_id not in self.friendly_platforms or enemy_id not in self.enemy_tracks:
            return None
        
        friendly = self.friendly_platforms[friendly_id]
        enemy_track = self.enemy_tracks[enemy_id]
        
        # 获取敌机状态估计
        enemy_state = enemy_track['estimator'].get_state()
        enemy_pos = enemy_state['position']
        enemy_vel = enemy_state['velocity']
        
        # 将我机位置转换为ENU坐标（以同一参考点）
        ref_platform = next(iter(self.friendly_platforms.values()))
        friendly_enu = CoordinateTransformer.llh_to_enu(
            friendly.latitude, friendly.longitude, friendly.altitude,
            ref_platform.latitude, ref_platform.longitude, ref_platform.altitude
        )
        
        # 计算我机速度矢量（简化处理）
        friendly_vel = self._estimate_friendly_velocity(friendly)
        
        # 计算时间间隔
        current_time = time.time()
        if enemy_track['time_history']:
            dt = current_time - enemy_track['time_history'][-1]
        else:
            dt = 1.0  # 默认1秒
        
        # 计算各威胁分量
        threat_components = self._calculate_threat_components(
            friendly_enu, friendly_vel, enemy_pos, enemy_vel, dt, enemy_track
        )
        
        # 计算动态权重
        dynamic_weights = self.calculate_dynamic_weights(friendly_id, enemy_id, current_time)
        
        # 计算综合威胁
        total_threat = 0
        for component, value in threat_components.items():
            if component in dynamic_weights:
                total_threat += dynamic_weights[component] * value
        
        # 记录威胁历史 
        threat_assessment = {
            'timestamp': current_time,
            'total_threat': total_threat,
            'components': threat_components,
            'weights': dynamic_weights
        }
        
        enemy_track['threat_history'].append(threat_assessment)
        
        return threat_assessment
    
    def _estimate_friendly_velocity(self, friendly: PlatParaInfo):
        """估计我机速度矢量"""
        # 根据速度、方位角、仰角计算速度矢量
        speed = friendly.speed
        
        # 将球坐标转换为直角坐标
        azimuth_rad = np.radians(friendly.azimuth)
        elevation_rad = np.radians(friendly.elevation)
        
        vx = speed * np.cos(elevation_rad) * np.sin(azimuth_rad)
        vy = speed * np.cos(elevation_rad) * np.cos(azimuth_rad)
        vz = speed * np.sin(elevation_rad)
        
        return np.array([vx, vy, vz])
    
    def _calculate_threat_components(self, friendly_pos, friendly_vel, enemy_pos, enemy_vel, dt, enemy_track):
        """计算各威胁分量"""
        components = {}
        
        # 1. 距离威胁
        distance = np.linalg.norm(enemy_pos - friendly_pos)
        components['distance'] = self._calculate_distance_threat(distance)
        
        # 2. 角度威胁
        components['angle'] = self._calculate_angle_threat(friendly_pos, enemy_pos, enemy_vel)
        
        # 3. 速度威胁
        v_closure = self.calculate_high_precision_closure_speed(
            friendly_pos, friendly_vel, enemy_pos, enemy_vel, dt
        )
        components['speed'] = self._calculate_speed_threat(v_closure)
        
        # 4. 行为威胁
        components['behavior'] = self._calculate_behavior_threat(enemy_track)
        
        return components
    
    def _calculate_distance_threat(self, distance):
        """计算距离威胁分量"""
        k_d = self.threat_params['distance']['k_d']
        d_max = self.threat_params['distance']['d_max']
        
        # 使用S型函数
        return 1 / (1 + np.exp(k_d * (distance - d_max * 0.3)))
    
    def _calculate_angle_threat(self, friendly_pos, enemy_pos, enemy_vel):
        """计算角度威胁分量"""
        if np.linalg.norm(enemy_vel) < 1e-6:
            return 0.0
        
        # 敌机指向我机的向量
        direction_to_friendly = friendly_pos - enemy_pos
        if np.linalg.norm(direction_to_friendly) < 1e-6:
            return 1.0
        
        direction_to_friendly_norm = direction_to_friendly / np.linalg.norm(direction_to_friendly)
        enemy_vel_norm = enemy_vel / np.linalg.norm(enemy_vel)
        
        # 计算夹角余弦值
        cos_theta = np.dot(enemy_vel_norm, direction_to_friendly_norm)
        
        # 确保在[0,1]范围内
        return max(cos_theta, 0)
    
    def _calculate_speed_threat(self, v_closure):
        """计算速度威胁分量"""
        v_max = self.threat_params['speed']['v_closure_max']
        return min(v_closure / v_max, 1.0)
    
    def _calculate_behavior_threat(self, enemy_track):
        """计算行为威胁分量"""
        if len(enemy_track['position_history']) < 3:
            return 0.0
        
        # 计算轨迹曲率
        curvature = self._calculate_trajectory_curvature(enemy_track['position_history'])
        
        # 计算机动强度
        maneuver_intensity = self._calculate_maneuver_intensity(enemy_track)
        
        # 综合行为威胁
        behavior_threat = 0.6 * min(curvature / self.threat_params['behavior']['curvature_threshold'], 1.0) + \
                         0.4 * min(maneuver_intensity / self.threat_params['behavior']['maneuver_threshold'], 1.0)
        
        return behavior_threat
    
    def _calculate_trajectory_curvature(self, positions):
        """计算轨迹曲率"""
        if len(positions) < 3:
            return 0.0
        
        positions_array = np.array(list(positions))
        
        # 使用三点法计算曲率
        curvatures = []
        for i in range(1, len(positions_array)-1):
            v1 = positions_array[i] - positions_array[i-1]
            v2 = positions_array[i+1] - positions_array[i]
            
            if np.linalg.norm(v1) < 1e-6 or np.linalg.norm(v2) < 1e-6:
                continue
                
            # 计算夹角
            cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
            angle = np.arccos(np.clip(cos_angle, -1, 1))
            curvatures.append(angle)
        
        return np.mean(curvatures) if curvatures else 0.0
    
    def _calculate_maneuver_intensity(self, enemy_track):
        """计算机动强度"""
        if len(enemy_track['position_history']) < 2:
            return 0.0
        
        positions = list(enemy_track['position_history'])
        accelerations = []
        
        # 计算逐点加速度
        for i in range(1, len(positions)-1):
            if len(enemy_track['time_history']) > i+1:
                dt = enemy_track['time_history'][i+1] - enemy_track['time_history'][i]
                if dt > 0:
                    vel1 = (positions[i] - positions[i-1]) / dt
                    vel2 = (positions[i+1] - positions[i]) / dt
                    acceleration = np.linalg.norm(vel2 - vel1) / dt
                    accelerations.append(acceleration)
        
        return np.mean(accelerations) if accelerations else 0.0
    
    def evaluate_all_threats(self):
        """评估所有敌机对所有我机的威胁"""
        threat_matrix = {}
        
        for friendly_id in self.friendly_platforms.keys():
            threat_matrix[friendly_id] = {}
            
            for enemy_id in self.enemy_tracks.keys():
                threat_assessment = self.evaluate_threat(friendly_id, enemy_id)
                if threat_assessment:
                    threat_matrix[friendly_id][enemy_id] = threat_assessment
        
        return threat_matrix
    
    def get_threat_trend(self, friendly_id, enemy_id, prediction_horizon=5):
        """获取威胁趋势预测"""
        if enemy_id not in self.enemy_tracks:
            return None
        
        track = self.enemy_tracks[enemy_id]
        threat_history = list(track['threat_history'])
        
        if len(threat_history) < 3:
            return None
        
        # 提取历史威胁值
        threats = [threat['total_threat'] for threat in threat_history[-5:]]
        
        # 简单线性回归预测
        x = np.arange(len(threats))
        y = np.array(threats)
        
        # 线性拟合
        A = np.vstack([x, np.ones(len(x))]).T
        slope, intercept = np.linalg.lstsq(A, y, rcond=None)[0]
        
        # 预测未来威胁
        future_threat = slope * (len(threats) + prediction_horizon) + intercept
        
        return {
            'current_threat': threats[-1] if threats else 0,
            'predicted_threat': max(0, min(future_threat, 1)),  # 限制在[0,1]范围
            'trend': 'increasing' if slope > 0.01 else 'decreasing' if slope < -0.01 else 'stable',
            'confidence': min(abs(slope) * 10, 1)  # 趋势置信度
        }

class RealTimeThreatAssessmentSystem:
    """实时威胁评估系统"""
    
    def __init__(self):
        self.threat_evaluator = TimeAwareThreatEvaluator()
        self.update_interval = 0.1  # 100ms更新间隔
        self.last_update_time = 0
        
    def process_real_time_data(self, friendly_data: List[PlatParaInfo], radar_data: List[RadarInfo]):
        """处理实时数据"""
        current_time = time.time()
        
        # 控制更新频率
        if current_time - self.last_update_time < self.update_interval:
            return None
            
        self.last_update_time = current_time
        
        # 更新我方信息
        for plat_info in friendly_data:
            self.threat_evaluator.update_friendly_info(plat_info)
        
        # 更新雷达信息
        self.threat_evaluator.update_radar_info(radar_data)
        
        # 评估威胁
        threat_matrix = self.threat_evaluator.evaluate_all_threats()
        
        return threat_matrix
    
    def get_comprehensive_threat_assessment(self):
        """获取综合威胁评估报告"""
        threat_matrix = self.threat_evaluator.evaluate_all_threats()
        
        report = {
            'timestamp': time.time(),
            'threat_assessments': {},
            'summary': {
                'max_threat': 0,
                'most_dangerous_enemy': None,
                'most_threatened_friendly': None
            }
        }
        
        max_threat = 0
        for friendly_id, enemies in threat_matrix.items():
            report['threat_assessments'][friendly_id] = {}
            
            for enemy_id, assessment in enemies.items():
                report['threat_assessments'][friendly_id][enemy_id] = assessment
                
                # 趋势预测
                trend = self.threat_evaluator.get_threat_trend(friendly_id, enemy_id)
                if trend:
                    report['threat_assessments'][friendly_id][enemy_id]['trend'] = trend
                
                # 更新最大威胁
                if assessment['total_threat'] > max_threat:
                    max_threat = assessment['total_threat']
                    report['summary']['max_threat'] = max_threat
                    report['summary']['most_dangerous_enemy'] = enemy_id
                    report['summary']['most_threatened_friendly'] = friendly_id
        
        return report

# 使用示例
def main():
    # 初始化系统
    threat_system = RealTimeThreatAssessmentSystem()
    
    # 模拟实时数据输入
    i=0.1
    while True:
        # 模拟我方数据
        friendly_data = [
            PlatParaInfo(
                id=1, longitude=116.3, latitude=39.9, altitude=5000,
                speed=200, roll=0, azimuth=45, elevation=10, timestamp=time.time()
            ),
            PlatParaInfo(
                id=2, longitude=116.4, latitude=39.8, altitude=5500,
                speed=180, roll=5, azimuth=30, elevation=5, timestamp=time.time()
            )
        ]
        
        # 模拟雷达数据
        radar_data = [
            RadarInfo(enemy_id=101, longitude=116.5+i, latitude=40.0+i, altitude=4800+i*2, timestamp=time.time()),

            RadarInfo(enemy_id=102, longitude=116.2+i, latitude=39.7+i, altitude=5200+i*2, timestamp=time.time())
        ]
        i+=i
        # 处理数据并获取威胁评估
        threat_matrix = threat_system.process_real_time_data(friendly_data, radar_data)
        
        if threat_matrix:
            print(f"\n=== 威胁评估报告 ===")
            print(f"时间: {time.strftime('%Y-%m-%d %H:%M:%S')}")
            
            for friendly_id, enemies in threat_matrix.items():
                print(f"\n我方平台 {friendly_id}:")
                for enemy_id, assessment in enemies.items():
                    print(f"  敌机 {enemy_id}: 威胁系数 = {assessment['total_threat']:.3f}")
                    print(f"    距离威胁: {assessment['components']['distance']:.3f}")
                    print(f"    角度威胁: {assessment['components']['angle']:.3f}")
                    print(f"    速度威胁: {assessment['components']['speed']:.3f}")
                    print(f"    行为威胁: {assessment['components']['behavior']:.3f}")
        
        # 获取详细报告（可选）
        if time.time() % 5 < 0.1:  # 每5秒输出一次详细报告
            detailed_report = threat_system.get_comprehensive_threat_assessment()
            print("\n=== 详细威胁分析 ===")
            print(f"最大威胁: {detailed_report['summary']['max_threat']:.3f}")
            print(f"最危险敌机: {detailed_report['summary']['most_dangerous_enemy']}")
            print(f"最受威胁平台: {detailed_report['summary']['most_threatened_friendly']}")
        
        time.sleep(0.1)  # 控制循环速度

if __name__ == "__main__":
    main()
